#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
bin_transfer_server.py

稳健的二进制分块传输服务器（与上面 lwIP 客户端配合）
- 明确小端字节序 '<'
- 使用 recv_exact() 避免 partial recv 导致错判
- 使用 '<IH' 格式打包块头（4B block_idx, 2B block_len）
- 每块等待客户端 5 字节响应（1B type + 4B idx）
- 新增：设备接入时询问是否升级功能
"""

import socket
import struct
import os
import threading
import time
from time import sleep

# CRC32 表（与 C 代码中 crc32_table 相同）
crc32_table = [
    0x00000000, 0x77073096, 0xee0e612c, 0x990951ba, 0x076dc419, 0x706af48f, 0xe963a535, 0x9e6495a3,
    0x0edb8832, 0x79dcb8a4, 0xe0d5e91e, 0x97d2d988, 0x09b64c2b, 0x7eb17cbd, 0xe7b82d07, 0x90bf1d91,
    0x1db71064, 0x6ab020f2, 0xf3b97148, 0x84be41de, 0x1adad47d, 0x6ddde4eb, 0xf4d4b551, 0x83d385c7,
    0x136c9856, 0x646ba8c0, 0xfd62f97a, 0x8a65c9ec, 0x14015c4f, 0x63066cd9, 0xfa0f3d63, 0x8d080df5,
    0x3b6e20c8, 0x4c69105e, 0xd56041e4, 0xa2677172, 0x3c03e4d1, 0x4b04d447, 0xd20d85fd, 0xa50ab56b,
    0x35b5a8fa, 0x42b2986c, 0xdbbbc9d6, 0xacbcf940, 0x32d86ce3, 0x45df5c75, 0xdcd60dcf, 0xabd13d59,
    0x26d930ac, 0x51de003a, 0xc8d75180, 0xbfd06116, 0x21b4f4b5, 0x56b3c423, 0xcfba9599, 0xb8bda50f,
    0x2802b89e, 0x5f058808, 0xc60cd9b2, 0xb10be924, 0x2f6f7c87, 0x58684c11, 0xc1611dab, 0xb6662d3d,
    0x76dc4190, 0x01db7106, 0x98d220bc, 0xefd5102a, 0x71b18589, 0x06b6b51f, 0x9fbfe4a5, 0xe8b8d433,
    0x7807c9a2, 0x0f00f934, 0x9609a88e, 0xe10e9818, 0x7f6a0dbb, 0x086d3d2d, 0x91646c97, 0xe6635c01,
    0x6b6b51f4, 0x1c6c6162, 0x856530d8, 0xf262004e, 0x6c0695ed, 0x1b01a57b, 0x8208f4c1, 0xf50fc457,
    0x65b0d9c6, 0x12b7e950, 0x8bbeb8ea, 0xfcb9887c, 0x62dd1ddf, 0x15da2d49, 0x8cd37cf3, 0xfbd44c65,
    0x4db26158, 0x3ab551ce, 0xa3bc0074, 0xd4bb30e2, 0x4adfa541, 0x3dd895d7, 0xa4d1c46d, 0xd3d6f4fb,
    0x4369e96a, 0x346ed9fc, 0xad678846, 0xda60b8d0, 0x44042d73, 0x33031de5, 0xaa0a4c5f, 0xdd0d7cc9,
    0x5005713c, 0x270241aa, 0xbe0b1010, 0xc90c2086, 0x5768b525, 0x206f85b3, 0xb966d409, 0xce61e49f,
    0x5edef90e, 0x29d9c998, 0xb0d09822, 0xc7d7a8b4, 0x59b33d17, 0x2eb40d81, 0xb7bd5c3b, 0xc0ba6cad,
    0xedb88320, 0x9abfb3b6, 0x03b6e20c, 0x74b1d29a, 0xead54739, 0x9dd277af, 0x04db2615, 0x73dc1683,
    0xe3630b12, 0x94643b84, 0x0d6d6a3e, 0x7a6a5aa8, 0xe40ecf0b, 0x9309ff9d, 0x0a00ae27, 0x7d079eb1,
    0xf00f9344, 0x8708a3d2, 0x1e01f268, 0x6906c2fe, 0xf762575d, 0x806567cb, 0x196c3671, 0x6e6b06e7,
    0xfed41b76, 0x89d32be0, 0x10da7a5a, 0x67dd4acc, 0xf9b9df6f, 0x8ebeeff9, 0x17b7be43, 0x60b08ed5,
    0xd6d6a3e8, 0xa1d1937e, 0x38d8c2c4, 0x4fdff252, 0xd1bb67f1, 0xa6bc5767, 0x3fb506dd, 0x48b2364b,
    0xd80d2bda, 0xaf0a1b4c, 0x36034af6, 0x41047a60, 0xdf60efc3, 0xa867df55, 0x316e8eef, 0x4669be79,
    0xcb61b38c, 0xbc66831a, 0x256fd2a0, 0x5268e236, 0xcc0c7795, 0xbb0b4703, 0x220216b9, 0x5505262f,
    0xc5ba3bbe, 0xb2bd0b28, 0x2bb45a92, 0x5cb36a04, 0xc2d7ffa7, 0xb5d0cf31, 0x2cd99e8b, 0x5bdeae1d,
    0x9b64c2b0, 0xec63f226, 0x756aa39c, 0x026d930a, 0x9c0906a9, 0xeb0e363f, 0x72076785, 0x05005713,
    0x95bf4a82, 0xe2b87a14, 0x7bb12bae, 0x0cb61b38, 0x92d28e9b, 0xe5d5be0d, 0x7cdcefb7, 0x0bdbdf21,
    0x86d3d2d4, 0xf1d4e242, 0x68ddb3f8, 0x1fda836e, 0x81be16cd, 0xf6b9265b, 0x6fb077e1, 0x18b74777,
    0x88085ae6, 0xff0f6a70, 0x66063bca, 0x11010b5c, 0x8f659eff, 0xf862ae69, 0x616bffd3, 0x166ccf45,
    0xa00ae278, 0xd70dd2ee, 0x4e048354, 0x3903b3c2, 0xa7672661, 0xd06016f7, 0x4969474d, 0x3e6e77db,
    0xaed16a4a, 0xd9d65adc, 0x40df0b66, 0x37d83bf0, 0xa9bcae53, 0xdebb9ec5, 0x47b2cf7f, 0x30b5ffe9,
    0xbdbdf21c, 0xcaebe28a, 0x53b3a330, 0x24b493a6, 0xbad02615, 0xcdd71683, 0x54de4739, 0x23d977af,
    0xb3666a2e, 0xc4615ab8, 0x5d680b02, 0x2a6f3b94, 0xb40bbe37, 0xc30c8ea1, 0x5a05df1b, 0x2d02ef8d
]


def crc32_incremental(data: bytes, prev_crc: int = 0) -> int:
    """
    Incrementally compute CRC32 using the same algorithm as in C code.

    Args:
        data: Data to calculate CRC on
        prev_crc: Previous CRC value (default is 0)

    Returns:
        Updated CRC32 value masked with 0xFFFFFFFF
    """
    crc = prev_crc ^ 0xFFFFFFFF
    for byte in data:
        crc = (crc >> 8) ^ crc32_table[(crc ^ byte) & 0xFF]
    return crc ^ 0xFFFFFFFF


class BinTransferServer:
    def __init__(self, host='0.0.0.0', port=1234, block_size=4096, max_retry=10, per_block_timeout=5.0):
        self.host = host
        self.port = port
        self.block_size = block_size
        self.max_retry = max_retry
        self.per_block_timeout = per_block_timeout  # 每块等待 ACK 超时（秒）
        self.bin_path = None
        self.file_size = 0
        self.total_blocks = 0
        self.global_crc32 = 0

    def set_bin_file(self, bin_path):
        """设置 bin 文件并预处理（大小、分块数、整体 CRC32）"""
        if not os.path.exists(bin_path):
            raise FileNotFoundError(f"文件不存在：{bin_path}")
        self.bin_path = bin_path
        self.file_size = os.path.getsize(bin_path)
        self.total_blocks = (self.file_size + self.block_size - 1) // self.block_size

        # 计算整体 CRC32 分段读取
        crc = 0
        with open(bin_path, 'rb') as f:
            while True:
                data = f.read(65536)
                if not data:
                    break
                crc = crc32_incremental(data, crc)
        self.global_crc32 = crc & 0xFFFFFFFF
        print(f"[INFO] 文件预处理完成：size={self.file_size}, blocks={self.total_blocks}, CRC32=0x{self.global_crc32:08X}")
        # 额外：用工具验证的提示
        print(f"[INFO] 建议用工具验证文件CRC32：python -c \"import zlib; print(hex(zlib.crc32(open('{bin_path}','rb').read()) & 0xFFFFFFFF))\"")

    def _calc_block_crc16(self, data: bytes) -> int:
        """计算 XMODEM CRC16（与 C 端实现匹配）"""
        crc = 0
        for b in data:
            crc ^= (b << 8)
            for _ in range(8):
                if crc & 0x8000:
                    crc = (crc << 1) ^ 0x1021
                else:
                    crc <<= 1
                crc &= 0xFFFF
        return crc

    def _send_file_header(self, conn: socket.socket):
        """
        发送文件头
        使用小端 '<'：<IIIII> —— magic(4), file_size(4), total_blocks(4), block_size(4), global_crc32(4)
        magic 使用 0x55AA55AA
        """
        header = struct.pack('<IIIII', 0x55AA55AA, self.file_size, self.total_blocks, self.block_size, self.global_crc32)
        conn.sendall(header)
        print("[INFO] 已发送文件头")

    def recv_exact(self, conn: socket.socket, nbytes: int, timeout: float) -> bytes | None:
        """
        循环读直到读满 nbytes，或在超时/对端关闭时返回 None。
        该函数会设置临时 timeout（函数返回后不会恢复之前的 timeout）。
        """
        orig_timeout = conn.gettimeout()
        try:
            conn.settimeout(timeout)
            buf = b''
            while len(buf) < nbytes:
                chunk = conn.recv(nbytes - len(buf))
                if not chunk:
                    # 对端关闭连接
                    return None
                buf += chunk
            return buf
        except socket.timeout:
            return None
        except Exception as e:
            print(f"[ERROR] recv_exact 异常: {e}")
            return None
        finally:
            try:
                conn.settimeout(orig_timeout)
            except Exception:
                pass

    def _handle_client(self, conn: socket.socket, addr):
        """处理单个客户端连接的核心逻辑"""
        print(f"[INFO] 客户端连接：{addr}")
        try:
            # 新增：询问用户是否升级
            print(f"\n[用户交互] 检测到设备 {addr} 接入，是否进行升级？")
            while True:
                user_choice = input("请输入 (1-确认升级，0-忽略并断开)：").strip()
                if user_choice in ['0', '1']:
                    break
                print("无效输入，请输入 1 或 0")

            if user_choice == '0':
                print(f"[INFO] 已选择忽略设备 {addr}，正在断开连接...")
                return  # 直接退出处理流程，会执行finally中的关闭连接操作

            # 减少延迟（禁用 Nagle）
            try:
                conn.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
            except Exception:
                pass

            # 1) 发送文件头
            self._send_file_header(conn)

            # 打开文件并按块发送
            with open(self.bin_path, 'rb') as f:
                # 初始化服务器端的块累积CRC32（与客户端逻辑一致）
                server_block_crc32 = 0  # 初始值与客户端相同（0）
                for block_idx in range(self.total_blocks):
                    f.seek(block_idx * self.block_size)
                    block_data = f.read(self.block_size)  # 读取当前块的原始文件数据
                    block_len = len(block_data)
                    block_crc16 = self._calc_block_crc16(block_data)

                    # 用当前块数据更新服务器端的累积CRC32，并打印
                    server_block_crc32 = crc32_incremental(block_data, server_block_crc32) & 0xFFFFFFFF
                    print(f"[SERVER] 块 {block_idx} 累积CRC32: 0x{server_block_crc32:08X} (数据长度: {block_len})")

                    # 组装包：<I H> 为 小端 4字节序号 + 2字节长度
                    block_header = struct.pack('<IH', block_idx, block_len)
                    packet = block_header + block_data + struct.pack('<H', block_crc16)

                    retry = 0
                    success = False
                    while retry < self.max_retry and not success:
                        # 发送数据块
                        try:
                            conn.sendall(packet)
                        except Exception as e:
                            raise Exception(f"发送块 {block_idx} 失败: {e}")
                        print(f"[SEND] 块 {block_idx}/{self.total_blocks} len={block_len} CRC=0x{block_crc16:04X} (尝试 {retry+1})")

                        # 等待客户端响应（5 字节）
                        resp = self.recv_exact(conn, 5, timeout=self.per_block_timeout)
                        if resp is None:
                            retry += 1
                            print(f"[WARN] 块 {block_idx} 响应超时/断开，重传 ({retry}/{self.max_retry})")
                            time.sleep(0.05)
                            continue

                        # 解析响应：1B type, 4B index，小端 '<BI'
                        try:
                            resp_type, resp_idx = struct.unpack('<BI', resp)
                        except struct.error:
                            retry += 1
                            print(f"[WARN] 块 {block_idx} 收到无效响应长度或格式，重传 ({retry}/{self.max_retry})")
                            continue

                        if resp_idx != block_idx:
                            print(f"[WARN] 响应序号不匹配：期望 {block_idx}，收到 {resp_idx}，忽略并重传")
                            retry += 1
                            continue

                        if resp_type == 0x00:
                            success = True
                            print(f"[ACK] 块 {block_idx} 确认接收")
                        else:
                            retry += 1
                            print(f"[NACK] 块 {block_idx} 校验失败，请求重传 ({retry}/{self.max_retry})")

                    if not success:
                        raise Exception(f"块 {block_idx} 超过最大重传次数，终止传输")

            # 发送完所有块，等待 客户端整体校验结果（1 字节）
            print("[INFO] 所有块发送完成，等待客户端整体校验结果...")
            final = self.recv_exact(conn, 1, timeout=10.0)
            if final is None:
                print("[WARN] 未收到客户端最终校验结果（超时或断开）")
            elif final == b'\x00':
                print("[INFO] 客户端整体校验成功，升级完成！")
            else:
                print("[ERROR] 客户端整体校验失败 (非 0x00)")

        except Exception as e:
            print(f"[ERROR] 传输过程中发生错误: {e}")
        finally:
            try:
                conn.close()
            except Exception:
                pass
            print(f"[INFO] 客户端连接已关闭：{addr}\n")

    def start(self):
        """启动服务器（阻塞，接收连接并为每个连接启动线程）"""
        if not self.bin_path:
            raise ValueError("请先调用 set_bin_file 设置 bin 路径")

        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
            s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            s.bind((self.host, self.port))
            s.listen(1)
            print(f"[INFO] 服务器启动，监听 {self.host}:{self.port}，等待客户端连接...")

            while True:
                conn, addr = s.accept()
                # 每个客户端用单独线程处理
                conn.settimeout(None)  # 我们在 recv_exact 中单独设置超时
                thr = threading.Thread(target=self._handle_client, args=(conn, addr), daemon=True)
                thr.start()


if __name__ == "__main__":
    # 配置
    BIN_FILE_PATH = "app1.bin"  # <- 改成你的 bin 文件路径
    SERVER_PORT = 1234
    BLOCK_SIZE = 4096                 # 与 MCU 一致
    MAX_RETRY = 30
    PER_BLOCK_TIMEOUT = 10.0           # 每块等待 ACK 的超时（秒）

    server = BinTransferServer(port=SERVER_PORT, block_size=BLOCK_SIZE, max_retry=MAX_RETRY, per_block_timeout=PER_BLOCK_TIMEOUT)
    try:
        server.set_bin_file(BIN_FILE_PATH)
        server.start()
    except Exception as ex:
        print(f"[FATAL] 启动失败: {ex}")